#include "EQPoisson.h"
#include "matrix_tool.h"
#include <cmath>
#include <vector>
#include <iostream>
#include <mgl2/mgl.h>

using namespace arma;
using namespace std;

void EQPoisson::set_xlimits(double lower, double upper)
{
	xlimits[0] = lower;
	xlimits[1] = upper;
}

void EQPoisson::set_ylimits(double lower, double upper)
{
	ylimits[0] = lower;
	ylimits[1] = upper;
}

bool EQPoisson::set_f(f2v f)
{
	if (f == NULL)
		return false;
	else
	{
		_f = f;
		return true;
	}
}

bool EQPoisson::set_BC_left(f1v BC_left)
{
	if (BC_left == NULL)
		return false;
	else
	{
		_BC_left = BC_left;
		return true;
	}  
}

bool EQPoisson::set_BC_right(f1v BC_right)
{
	if (BC_right == NULL)
		return false;
	else
	{
		_BC_right = BC_right;
		return true;
	}	
}

bool EQPoisson::set_BC_top(f1v BC_top)
{
	if (BC_top == NULL)
		return false;
	else
	{
		_BC_top = BC_top;
		return true;
	}	
}

bool EQPoisson::set_BC_bottom(f1v BC_bottom)
{
	if (BC_bottom == NULL)
		return false;
	else
	{
		_BC_bottom= BC_bottom;
		return true;
	}
}

void EQPoisson::set_Nx(int myNx)
{
	Nx = myNx;
}

void EQPoisson::set_Ny(int myNy)
{
	Ny = myNy;
}

void EQPoisson::set_eps_x(double epsx)
{
	eps_x = epsx;
}

void EQPoisson::set_eps_y(double epsy)
{
	eps_y = epsy;
}

void EQPoisson::set_eps(double x, double y)
{
	eps_x = x;
	eps_y = y;
}

double EQPoisson::determine_hx(int method)
{
	switch(method)
	{
		case GAUSS_SEIDEL:
			hx = (xlimits[1] - xlimits[0])/Nx;
			return hx;
		case BLOCK_CYCLE_REDUCE:
			hx = (xlimits[1] - xlimits[0])/Nx;
			return hx;
	}
	
}

double EQPoisson::determine_hy(int method)
{	
	switch(method)
	{
		case GAUSS_SEIDEL:
			hy = (ylimits[1] - ylimits[0])/Ny;
			return hy;
		case BLOCK_CYCLE_REDUCE:
			hy = (ylimits[1] - ylimits[0])/Ny;
			return hy;
	}
}

unsigned int EQPoisson::compute_Nx(int method)
{
	switch(method)
	{
		case GAUSS_SEIDEL:
			return 0;
		case BLOCK_CYCLE_REDUCE:
			int k = ceil(log((xlimits[1] - xlimits[0])/eps_x)/log(2));
			Nx = pow(2,k);
			return Nx;
	}
	
}

unsigned int EQPoisson::compute_Ny(int method)
{
	switch(method)
	{
		case GAUSS_SEIDEL:
			return 0;
		case BLOCK_CYCLE_REDUCE:
			int k = ceil(log((ylimits[1] - ylimits[0])/eps_y)/log(2));
			_k = k;
			Ny = pow(2,k);
	}

}

void EQPoisson::init(int method)
{
	compute_Nx(method);
	determine_hx(method);
	compute_Ny(method);
	determine_hy(method);
	//-----------------------------------------------------------
	cout << "x in [" << xlimits[0] << ", " << xlimits[1] << "]\n";
	cout << "y in [" << ylimits[0] << ", " << ylimits[1] << "]\n";
	if (method == BLOCK_CYCLE_REDUCE)
	{
		cout << "k = " << _k << endl;
	}
	cout << "Nx = " << Nx << ", Ny = " << Ny << endl;
	cout << "hx = " << hx << ", hy = " << hy << endl;
	//-----------------------------------------------------------
	u.set_size((Nx-1),(Ny-1));
	S.set_size(Nx-1,Nx-1);
	S = S.eye()/hy/hy;
	// --------------------------------------------------------
	// cout << "u.size()=" << u.size() << endl;
	// --------------------------------------------------------
	s = 1/(hx*hx);
	d = -2*(hx*hx + hy*hy)/(hx*hx *hy*hy);
	D.zeros(Nx-1,Nx-1);
	D(0,0) = d; D(0,1) = s;
	
	for (int i = 1; i < Nx-2; ++i)
	{
		D(i,i-1) = s;
		D(i,i) = d;
		D(i,i+1) = s;
	}
	
	D(Nx-2,Nx-2) = d; D(Nx-2,Nx-3) = s;
	// ---------------------------------------------------
	// cout << "D.size()=" << D.size() << endl;
	// -----------------------------------------------------
	b.set_size((Nx-1)*(Ny-1));
	
	b(1-1) = (*_f)(xlimits[0]+hx,ylimits[0]+hy)-(*_BC_left)(xlimits[0]+hx)/hy/hy-(*_BC_top)(ylimits[0]+hy)/hx/hx;
	
	for (int i = 2; i <= Nx-2; ++i)
	{
		b(i-1)=(*_f)(xlimits[0]+hx*i,ylimits[0]+hy) - (*_BC_left)(xlimits[0]+hx*i)/hy/hy;
	}

	b(Nx-1-1) = (*_f)(xlimits[0]+hx*(Nx-1),ylimits[0]+hy)-(*_BC_left)(xlimits[0]+hx*(Nx-1))/hy/hy-(*_BC_bottom)(ylimits[0]+hy)/hx/hx;
	
	for (int i = 2; i <= Ny-2; ++i)
	{
		b((Nx-1)*(i-1)+1-1) = (*_f)(xlimits[0]+hx,ylimits[0]+hy*i) - (*_BC_top)(ylimits[0]+i*hy)/hx/hx;
		
		for (int j = 2; j <= Nx-2; ++j)
		{
			b(i*(Nx-1)+j-1) = (*_f)(xlimits[0]+j*hx,ylimits[0]+hy*i);
		}
		
		b((Nx-1)*i-1) = (*_f)(xlimits[0]+(Nx-1)*hx,ylimits[0]+hy*i) - (*_BC_bottom)(ylimits[0]+i*hy)/hx/hx;
	}

	b((Nx-1)*(Ny-2)+1-1) = (*_f)(xlimits[0]+hx,ylimits[0]+hy*(Ny-1))-(*_BC_right)(xlimits[0]+hx)/hy/hy-(*_BC_top)(ylimits[0]+(Ny-1)*hy)/hx/hx;
	
	for (int i = 2; i <= Nx-2; ++i)
	{
		b((Nx-1)*(Ny-2)+i-1)=(*_f)(xlimits[0]+hx*i,ylimits[0]+hy*(Ny-1)) - (*_BC_right)(xlimits[0]+hx*i)/hy/hy;
	}

	b((Nx-1)*(Ny-1)-1)= (*_f)(xlimits[0]+hx*(Nx-1),ylimits[0]+hy*(Ny-1))-(*_BC_right)(xlimits[0]+hx*(Nx-1))/hy/hy-(*_BC_bottom)(ylimits[0]+(Ny-1)*hy)/hx/hx;
	// ----------------------------------------------------------------
	// cout << "b.size()=" << b.size() << endl;

	// (b*hx*hx).print("b=");
	// ------------------------------------------------------------------------
}

void EQPoisson::block_cycle_reduce_method()
{
	vector<mat> iter_D(_k);
	vector<mat> iter_S(_k);
	vector<mat> iter_b(_k);
	unsigned int r;

	iter_S[0] = S*(hx*hy);
	iter_D[0] = D*(hx*hy);

	iter_b[0].set_size(Nx-1,Ny-1);
	
	for (int i = 0; i < Ny-1; ++i)
	{
		iter_b[0].col(i) = b.subvec(i*(Nx-1),(i+1)*(Nx-1)-1)*hx*hy;
	}
	// D.print("D");
	// S.print("S");
	// b.print("b");
	for (int i = 1; i <= _k-1; ++i)
	{
		iter_b[i].set_size(Nx-1,Ny-1);
		iter_S[i] = iter_S[i-1]*iter_S[i-1];
		iter_D[i] = 2*iter_S[i] - iter_D[i-1]*iter_D[i-1];
		r = (unsigned int) pow(2,i);

		for (int j = 1; j <= (int) pow(2,_k-i)-1; ++j)
		{
			iter_b[i].col(j*r-1) = iter_S[i-1]*(iter_b[i-1].col(j*r-r/2-1)+iter_b[i-1].col(j*r+r/2-1)) - iter_D[i-1]*iter_b[i-1].col(j*r-1);
		}
	}
	// cout << "problem?------------No!\n";
	vector<vec> x(Ny-1);
	for (int i = 0; i < Ny-1; ++i)
	{
	 	x[i].set_size(Nx-1);
	} 

	x[(int) pow(2,_k-1)-1] = iter_D[_k-1].i()*iter_b[_k-1].col((int) pow(2,_k-1)-1);
	// x[(int) pow(2,_k-1)-1].print("the last column");
	for (int i = _k-2; i >= 0; --i)
	{
		r = (unsigned int) pow(2,i);
		for (int j = 1; j <= (int) pow(2,_k-i-1); ++j)
		{
			vec c(Nx-1); 
			if (1==j)
			{
				c = iter_b[i].col((2*j-1)*r-1) - iter_S[i]*x[2*j*r-1];
			}
			else if ((int) pow(2,_k-i-1) == j)
					c = iter_b[i].col((2*j-1)*r-1) - iter_S[i]*x[(2*j-2)*r-1];
				else
					c = iter_b[i].col((2*j-1)*r-1) - iter_S[i]*(x[2*j*r-1]+x[(2*j-2)*r-1]);
			x[(2*j-1)*r-1] = iter_D[i].i()*c;
		}
		
	}
	for (int i = 0; i < Ny-1; ++i)
	{
		u.col(i) = x[i];
	}

	vec u_left(Nx-1); 
	vec u_right(Nx-1);
	rowvec u_top(Ny+1);
	rowvec u_bottom(Ny+1);
	for (int i = 0; i < Nx-1; ++i)
	{
		u_left(i) = (*_BC_left)(xlimits[0]+hx*(i+1));
		u_right(i) = (*_BC_right)(xlimits[0]+hx*(i+1));
	}
	for (int i = 0; i < Ny+1; ++i)
	{
		u_top(i) = (*_BC_top)(xlimits[0]+hx*i);
		u_bottom(i) = (*_BC_bottom)(xlimits[0]+hx*i);
	}

	u.insert_cols(0,u_left);
	u.insert_cols(Ny,u_right);
	// u.swap_cols(Ny-1,Ny);
	u.insert_rows(0,u_top);
	u.insert_rows(Nx,u_bottom);
	// u.swap_rows(Nx-1,Nx);

}

/*double compute_A(int i,int j)
{
	if (i == j)
		return d;
	if (1 == i-j || 1 == j-i)
		return 1/hx/hx;
	if (Nx-1 == i-j || Nx-1 == j-i)
		return s;
	return 0.0;
}*/

/*double EQPoisson::GS_iterate(int i, vec x)
{
	double x;
	for (int j = 0; j <= i-2; ++j)
	{
		x -= compute_A(i,j)*x(j);
	}
	for (int j = i; j <= length; ++j)
	{
		x -= compute_A(i,j)*x(j);
	}

}*/


void EQPoisson::Gauss_Seidel_method(double precision=1e-4)
{
	int length = (Nx-1)*(Ny-1);
	vec x_new = ones<vec>(length);
	vec x_old;
	rowvec U(length); U.zeros();
	vec D_L(length);  D_L.zeros();
	U(1) = -s; U(Nx-1) = -1/hy/hy;
	D_L(0) = d;
	D_L(1) = s; D_L(Nx-1) = 1/hy/hy;
	vec D_L_inv = toeplitz_inverse(D_L);
	vec D_L_inv_b = col_toeplitz_mul(D_L_inv,b);
	
	// D_L_inv_b.print("D_L_inv_b");
	// D.print("D=");
	// D_L.print("D_L=");
	// D_L_inv.print("D_L.inv=");
	// U.print("U=");
	int count1 = 0;
	do
	{
		++count1;
		x_old = x_new;
		x_new = col_toeplitz_mul(D_L_inv,row_toeplitz_mul(U,x_old)) + D_L_inv_b;
		// x = row_toeplitz_mul(U,x);
		//x_old.print("x");
		// cout << "-----------------------------------\n";
		// row_toeplitz_mul(U,x[0]).print("Ux");
		// col_toeplitz_mul(D_L_inv,row_toeplitz_mul(U,x[0])).print("first");
		// col_toeplitz_mul(D_L_inv,b).print("second");
		// std::cout << arma::abs(x[0]-x[1]).max() << std::endl;
		// std::cout << arma::abs(x_new-x_old).max() << "\t" << count1 << std::endl;
	}
	while(arma::abs(x_new-x_old).max() > precision);
	// while(--count);

	for (int i = 0; i < Ny-1; ++i)
	{
		u.col(i) = x_new.rows(i*(Nx-1),(i+1)*(Nx-1)-1);
	}

	vec u_left(Nx-1); 
	vec u_right(Nx-1);
	rowvec u_top(Ny+1);
	rowvec u_bottom(Ny+1);
	for (int i = 0; i < Nx-1; ++i)
	{
		u_left(i) = (*_BC_left)(xlimits[0]+hx*(i+1));
		u_right(i) = (*_BC_right)(xlimits[0]+hx*(i+1));
	}
	for (int i = 0; i < Ny+1; ++i)
	{
		u_top(i) = (*_BC_top)(xlimits[0]+hx*i);
		u_bottom(i) = (*_BC_bottom)(xlimits[0]+hx*i);
	}

	u.insert_cols(0,u_left);
	u.insert_cols(Ny,u_right);
	// u.swap_cols(Ny-1,Ny);
	u.insert_rows(0,u_top);
	u.insert_rows(Nx,u_bottom);
	// u.swap_rows(Nx-1,Nx);
}

bool EQPoisson::solve(int method)
{
	init(method);
	
	switch (method)
	{
		case GAUSS_SEIDEL:
			Gauss_Seidel_method();
			break;

		case BLOCK_CYCLE_REDUCE:
			block_cycle_reduce_method();
			break;
	}

	if (0 == u.size())
		return false;
	else 
		return true;
}

void EQPoisson::save_data(const char* filename) const
{
	using namespace std;

	ofstream fout;
	fout.open(filename);
	//fout.precision(4);
	//fout.width(5);
	
	for (int i = 0; i < Nx+1; ++i)
	{
		for (int j = 0; j < Ny; ++j)
		{
			fout << u(i,j) << ", ";			
		}
		fout << u(i,Ny) << ";" << endl;
	}

}

void EQPoisson::draw_data(const char* filename) const
{
	mglData data(Nx+1, Ny+1);
	for (int i = 0; i < Nx+1; ++i)
	{
		
		for (int j = 0; j < Ny+1; ++j)
			data.a[i + j*(Nx+1)] = u(i,j);
	}

	mglGraph gr;
	gr.Light(true); 
	gr.Rotate(20,-30); 
	gr.Box();
	gr.SetRanges(xlimits[0],xlimits[1],ylimits[0],ylimits[1],floor(u.min()),ceil(u.max()));
	gr.Axis();
	gr.Grid();
	gr.Label('x',"x",1);
	gr.Label('y',"y",1);
	gr.Label('z',"u",1);
	gr.Surf(data);
	gr.WriteFrame(filename);
}